import argparse
from typing import Dict, List
import re
import json

from .Tools.CodeParser import CodeParser
from .Tools.TreeVisualizer import TreeVisualizer
from .Tools.CallGraphAnalyzer import CallGraphAnalyzer


import os
cwd = os.getcwd()
os.environ['PYTHONPATH'] = cwd

def analyze_python_code(
    src_directory: List[str], 
    target_function: str,
    is_recursive: bool = True,
    local_only: bool = True
):
    """
    Analyze Python code and generate call trees, chains, and computation steps.
    
    Args:
        src_directory: List of directories containing Python source files
        target_function: Target function name for specific analysis
        is_recursive: Whether to recursively scan directories
        local_only: Whether to only analyze modules within src_directory
    """
    # Parse the code
    parser = CodeParser(src_directory, is_recursive, local_only=local_only)
    print(f"Scanning directories: {src_directory}")
    # Turn on debug mode to help diagnose wildcard import issues
    parser.debug_mode = True
    parser.analyze_main_blocks()
    
    # Create analyzer from parsed data
    analyzer = CallGraphAnalyzer(parser)
    
    print(f"Found {len(parser.functions)} functions in {len(parser.modules)} modules")
    
    # Special handling for wildcard imports - make sure all imported functions are available
    for module_name in parser.modules:
        # If this module has wildcard imports
        imports = parser.get_module_imports(module_name)
        if '*' in imports:
            wildcard_modules = imports['*']
            print(f"Processing wildcard imports from {module_name}: {wildcard_modules}")
            # Process each wildcard import to add all functions to the call graph
            for imported_module in wildcard_modules:
                parser.process_wildcard_module_dependencies(module_name, imported_module)
    
    # Generate call tree for the entire program
    down_tree = analyzer.build_call_tree(direction="down", max_depth=10)
    # if target_function has appendix ".py", remove it
    if target_function.endswith(".py"):
        target_function = target_function[:-3]
    # Build a focused tree that only starts from main module
    main_module = target_function + ".__main_block__"
    filtered_tree = {}
    
    # Only include the branch starting from main_module
    if main_module in down_tree:
        filtered_tree[main_module] = down_tree[main_module]
    else:
        print(f"Warning: Main module '{main_module}' not found in call tree.")
        # Create an empty placeholder if main module isn't found
        filtered_tree[main_module] = {}
    
    print(f"\nCall tree for main execution path:")
    print(f"Tree has {len(filtered_tree)} root nodes")
    
    # Create Output folder if it doesn't exist
    output_dir = "Output"
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # Save both tree visualizations
    TreeVisualizer.visualize_call_tree(filtered_tree, "Output/tree_full_program.txt")
    
    # Show some statistics about the program
    print("\nProgram structure overview:")
    module_counts = {}
    for func_name in parser.functions:
        module_name = func_name.split('.')[0] if '.' in func_name else 'unknown'
        module_counts[module_name] = module_counts.get(module_name, 0) + 1
    
    for module, count in module_counts.items():
        print(f"  Module {module}: {count} functions")
    
    # Extract computation steps for all functions
    steps = analyzer.extract_computation_steps()
    TreeVisualizer.visualize_computation_steps(steps, "Output/steps_full_program.txt")
    # Extract function call chains
    chains = analyzer.generate_call_chains()
    TreeVisualizer.visualize_call_chains(chains, "Output/chains_full_program.txt")

def main():
    """Entry point for command-line usage"""
    src_default = ['.'] # default test directory
    target_default = "python_flow_map.py"
    parser = argparse.ArgumentParser(description="Python Flow Map: Analyze Python code structure and call relationships")
    parser.add_argument('--src', nargs='+', default=src_default, help='Source directories to scan')
    parser.add_argument('--target', default=target_default, help='Target function name for specific analysis')
    parser.add_argument('--recursive', action='store_true', default=True, help='Recursively scan directories')
    parser.add_argument('--local-only', action='store_true', default=True, 
                        help='Only analyze modules within source directories')
    
    args = parser.parse_args()
    
    analyze_python_code(
        args.src,
        args.target,
        args.recursive,
        args.local_only
    )

# For direct module execution
if __name__ == "__main__":
    main()
